import os
import torch
import csv
from torchvision import transforms
from PIL import Image
import argparse
from torchvision.transforms.functional import InterpolationMode

from dataset.augmentation.classification_preset_eval import ClassificationPresetEval


# Function to load model
def get_class_by_name(module_name, class_name):
    module = __import__(module_name, fromlist=[class_name])
    return getattr(module, class_name)


# Function to dynamically import a model class from a string name
def get_model_class_by_name(module_base, class_name):
    module_name = f'{module_base}.{class_name}'
    module = __import__(module_name, fromlist=[class_name])
    return getattr(module, class_name)


# Function to predict and save results
def predict_and_save(model, data_dir, output_csv):
    # 使用与验证集完全一致的预处理方式
    transform = ClassificationPresetEval(
        crop_size=224,
        resize_size=256,
        mean=(0.485, 0.456, 0.406),
        std=(0.229, 0.224, 0.225),
        interpolation=InterpolationMode.BILINEAR
    )

    with open(output_csv, mode='w', newline='', encoding='utf-8') as file:
        writer = csv.writer(file)
        for img_name in os.listdir(data_dir):
            img_path = os.path.join(data_dir, img_name)
            try:
                image = Image.open(img_path).convert('RGB')
                image = transform(image)  # 使用ClassificationPresetEval进行处理
                image = image.unsqueeze(0)  # Add batch dimension
                image = image.to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))

                with torch.no_grad():
                    outputs = model(image)
                    _, predicted = torch.max(outputs, 1)
                    label = f"{predicted.item():04d}"

                writer.writerow([img_name, label])
            except Exception as e:
                print(f"处理图片 {img_name} 时出错: {str(e)}")


# Main function to parse arguments and run prediction
def parse_args():
    parser = argparse.ArgumentParser(description='Configuration for model prediction')

    # Data settings
    parser.add_argument('--data_type', type=int, default=1,
                        help='test data type (1 for webfg400, other for webinat5000)')

    # Model settings
    parser.add_argument('--model_name', type=str, default='SwinTransformer',
                        help='选择模型的类名，例如：resnet101')
    parser.add_argument('--weight_path', type=str,
                        default=r'D:\competition\2026quanqiu\code\weights\best_accuracy_epoch144_acc43.01.pth',
                        help='Path to the trained model weights (.pth file)')
    parser.add_argument('--num_classes', type=int, default=400,
                        help='Number of classes')
    # Other settings
    parser.add_argument('--random_seed', type=int, default=42,
                        help='Random seed for reproducibility')

    return parser.parse_args()


def load_model(args, device):
    """Load the trained model"""
    # Select model
    model_class = get_model_class_by_name('model.backbone', args.model_name)
    model = model_class(num_classes=args.num_classes)

    # Load trained weights
    try:
        state_dict = torch.load(args.weight_path, map_location=device)
        model.load_state_dict(state_dict, strict=True)
        print(f"✓ Successfully loaded model weights from {args.weight_path}")
    except Exception as e:
        print(f"✗ Error loading model weights from {args.weight_path}: {str(e)}")
        raise e

    # Move model to device and set to evaluation mode
    model.to(device)
    model.eval()

    return model


def main(args):
    # 设置随机种子确保可重复性
    torch.manual_seed(args.random_seed)

    # 根据数据类型选择数据目录和输出文件
    if args.data_type == 1:
        data_dir = r'D:\competition\2026quanqiu\data\webfg400_test_A\test_A'
        output_csv = r'D:\competition\2026quanqiu\cons\w400\pred_results_web400.csv'
    else:
        data_dir = r'D:\competition\2026quanqiu\data\webinat5000_test_A\test_A'
        output_csv = r'D:\competition\2026quanqiu\cons\w5000\pred_results_web5000.csv'

    # 确保输出目录存在
    os.makedirs(os.path.dirname(output_csv), exist_ok=True)

    # 加载模型
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"使用设备: {device}")
    model = load_model(args, device)

    # 预测并保存结果
    predict_and_save(model, data_dir, output_csv)
    print(f"预测完成，结果已保存至 {output_csv}")


if __name__ == "__main__":
    args = parse_args()
    main(args)
